import torch
from train_and_save_model import MLP
import gradio as gr
import torchvision.transforms as transforms
import torchvision.datasets as datasets

# 从硬盘加载保存的模型
model = MLP()
model_path = './mnist_mlp.pth'
model.load_state_dict(torch.load(model_path))
model.eval()

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

#定义预测函数，这个函数将用于Gradio接口进行预测
def predict(image): 
    if image is None:  # 如果输入为空，则返回None
        return None 
    image = transform(image).unsqueeze(0)   # 转换图像数据并增加一个维度
    output = model(image)           # 输入模型得到预测结果
    _, predicted = torch.max(output.data, 1)    # 获取预测概率最大的类别索引
    return predicted.item()    
 
# 创建Gradio接口，这个接口将用于用户输入和显示预测结果
interface = gr.Interface(predict,
                     inputs=gr.Sketchpad(label="image"),
                     outputs=gr.Label(label="output"),
                     title="手写数字识别",  
                     description="使用pytorch构建的MLP模型对手写数字进行识别。",
                     live=False) 

# 启动Gradio接口，用户可以通过这个接口进行交互
interface.launch(share=True)